import copy
import os
import math
import statistics
from sklearn.cluster import KMeans
import open3d as o3d
import torch
import pickle as pl
# import mmcv
from scipy.spatial import Delaunay, distance
import numpy as np
import struct
# from mmdet3d.core.bbox.box_np_ops import *

def pc_cluster(points, n_clusters):
    kmeans = KMeans(n_clusters=n_clusters, random_state=0, n_init="auto").fit(points)
    labels = kmeans.labels_
    centers = kmeans.cluster_centers_
    return labels, centers


def filter_labels(labels, centers):
    # idx = [i for i in range(len(centers))]
    label = np.argwhere(centers[...,-1] <= min(centers)[-1])
    mask = np.array([0] * len(labels))
    valid_idx = np.argwhere(labels == label).reshape(-1)
    mask[valid_idx] = 1
    placeable_idx = labels[mask]
    return placeable_idx


def display_inlier_outlier(cloud, ind):
    inlier_cloud = cloud.select_by_mask(ind)
    outlier_cloud = cloud.select_by_mask(ind, invert=True)  # invert利用索引反选

    print("Showing outliers (red) and inliers (gray): ")
    outlier_cloud.paint_uniform_color([1, 0, 0])
    inlier_cloud.paint_uniform_color([0.8, 0.8, 0.8])
    o3d.visualization.draw([inlier_cloud, outlier_cloud])


def voxel_grid(points_3d, voxel_size, pc_range, min_height):
    # points_3d = point_cloud[:, :-1]
    pc_valid_mask = np.logical_and.reduce((points_3d >= pc_range[:3], points_3d < pc_range[3:]), axis=0)
    pc_valid_idx = np.argwhere(np.array(pc_valid_mask).sum(axis=1) == 3).reshape(-1)
    points_3d = np.array(points_3d)[pc_valid_idx,:]
    # point_cloud = np.array(point_cloud)[pc_valid_idx, :]
    min_bound = np.array(pc_range[:3])
    voxel_coords = np.floor((points_3d - min_bound) / voxel_size).astype(np.int32)
    voxel_count = np.max(voxel_coords, axis=0) + 1

    # 初始化每个体素块的高度列表
    voxel_heights = [[] for _ in range(np.prod(voxel_count))]
    # voxel_heights = [[] * np.prod(voxel_count)]

    # 将每个点添加到对应的体素块中
    for i, coord in enumerate(voxel_coords):
        index = np.ravel_multi_index(coord, voxel_count)
        voxel_heights[index].append(points_3d[i, 2])

    # 计算每个体素块的高度均值
    voxel_height_means = np.zeros(np.prod(voxel_count))
    voxel_height_max = np.zeros(np.prod(voxel_count))
    voxel_height_quantiles = np.zeros(np.prod(voxel_count))
    valid_idx, rows, cols = [], [], []
    for i, heights in enumerate(voxel_heights):
        if len(heights) > 1:
            voxel_height_max[i] = np.max(heights)
            # voxel_height_means[i] = np.mean(heights)
            # voxel_height_quantiles[i] = statistics.quantiles(heights, n=100)[-2]
            valid_idx.append(i)
            cols.append(i % voxel_count[1])
            rows.append(i // voxel_count[1])
    valid_coors = np.c_[rows, cols]
    # ground_mask_mean = voxel_height_means[valid_idx] < min_height
    # ground_mask_quantiles = voxel_height_quantiles[valid_idx] < min_height
    # ground_mask = ground_mask_mean & ground_mask_quantiles
    ground_mask_max = voxel_height_max[valid_idx] < min_height
    ground_mask = ground_mask_max
    
    valid_coors = valid_coors[ground_mask,:]
    valid_coors = np.c_[valid_coors, np.zeros(len(valid_coors))]
    return valid_coors, voxel_coords, points_3d


def noise_filter(points):
    voxel_down_pcd = points.voxel_down_sample(voxel_size=0.2)
    # o3d.visualization.draw([voxel_down_pcd])
    cl, ind = voxel_down_pcd.remove_statistical_outliers(nb_neighbors=10, std_ratio=3)
    # display_inlier_outlier(voxel_down_pcd, ind)
    points = voxel_down_pcd.select_by_mask(ind)
    return points

def select_placeable_area(point_cloud, voxel_coors, valid_coors, center_constraints = None):
    valid_idx = []
    for i in range(len(valid_coors)):
        coors = valid_coors[i, :]
        idx = np.where((voxel_coors == coors).all(1))[0]
        valid_idx.extend(idx)
    placeable_point_cloud = point_cloud[valid_idx]
    device = o3d.core.Device("CPU:0")
    dtype = o3d.core.float32
    pcd = o3d.t.geometry.PointCloud(o3d.core.Tensor(placeable_point_cloud[:, :3], dtype, device))
    placeable_point_cloud = noise_filter(pcd)
    placeable_point_cloud = placeable_point_cloud.point['positions'].numpy()

    if center_constraints is not None:
        l, w, h = center_constraints
        locked_min_bound = [-l/2, -w/2, -h/2]
        locked_max_bound = [l/2, w/2, h/2]
        invalid_mask = np.logical_and.reduce((placeable_point_cloud >= locked_min_bound, placeable_point_cloud <= locked_max_bound), axis=0)
        invalid_idx = np.argwhere(np.array(invalid_mask).sum(axis=1) == 3).reshape(-1)
        placeable_point_cloud = np.delete(placeable_point_cloud, invalid_idx,0)
    return placeable_point_cloud


def generate_road_plane(points):
    device = o3d.core.Device("CPU:0")
    dtype = o3d.core.float32
    pcd = o3d.t.geometry.PointCloud(o3d.core.Tensor(points[:, :3], dtype, device))
    points = noise_filter(pcd)
    plane_model, inliers = pcd.segment_plane(distance_threshold=0.1, ransac_n=10, num_iterations=1000)
    return plane_model


def placeable_area_filter(points, n = 100, min_height=0.01, pc_range=None):
    # points = np.fromfile(pc_file, dtype=np.float32)
    # points = points.reshape((-1,4))
    # points = np.fromfile(file_path).reshape((-1,4))
    device = o3d.core.Device("CPU:0")
    dtype = o3d.core.float32
    pcd = o3d.t.geometry.PointCloud(o3d.core.Tensor(points[:, :3], dtype, device))
    # pcd = o3d.geometry.PointCloud()
    # pcd.point['positions'] = o3d.core.Tensor(points[:, :3], dtype, device)
    # pcd.point['intensity'] = o3d.core.Tensor(points[:, 3], dtype, device)
    # pcd = o3d.t.io.read_point_cloud(pc_file)
    # noise filter
    points = noise_filter(pcd)
    # compute mean height of each pillar
    point_cloud = points.point['positions'].numpy()
    # point_intensity = points.point['intensity'].numpy()
    # point_cloud = np.c_[point_3d, point_intensity]
    pc_range = np.array(pc_range)
    grid_size = np.array([n, n, 1])
    voxel_size = (pc_range[3:] - pc_range[:3]) / grid_size
    valid_coors , voxel_coors, point_cloud = voxel_grid(point_cloud, voxel_size, pc_range, min_height)
    # 选择可放置点云位置
    placeable_pc = select_placeable_area(point_cloud, voxel_coors, valid_coors, center_constraints = None)
    return placeable_pc





def select_and_generate_best_bbox(valid_aug_bboxes, gt_bbox, pc):
    best_idx = np.argmax(np.cos(valid_aug_bboxes[:, -1] - gt_bbox[-1]))
    best = valid_aug_bboxes[best_idx]
    ground_mean_height = np.mean(pc[:, 2])
    z = ground_mean_height + gt_bbox[2] / 2
    aug_bbox = np.zeros(gt_bbox.shape[-1])
    aug_bbox[:2] = best[:2]
    aug_bbox[2] = z
    aug_bbox[3:6] = gt_bbox[3:6]
    aug_bbox[-1] = best[-1]
    return aug_bbox

def translate_points(box, pc_box):
    location = box[:, :3]
    pc_box = torch.Tensor(pc_box)
    location = torch.Tensor(location)
    pc_box[:, :3] += location
    pc_box = pc_box.numpy()
    return pc_box


def bin_dump(point_cloud, file_path):
    point_cloud.tofile(file_path)


def bin2pcd(bin_file):
    pcd_file = bin_file.replace('.bin', '.pcd')
    points = np.fromfile(bin_file, dtype="float32").reshape((-1, 4))
    size_float = 4
    list_pcd = []
    with open(bin_file, "rb") as f:
        byte = f.read(size_float * 4)
        while byte:
            x, y, z, intensity = struct.unpack("ffff", byte)
            list_pcd.append([x, y, z, intensity])
            byte = f.read(size_float * 4)
    np_pcd = np.asarray(list_pcd)
    # pcd = o3d.geometry.PointCloud()
    # pcd.points = o3d.utility.Vector3dVector(np_pcd)
    device = o3d.core.Device("CPU:0")
    dtype = o3d.core.float32
    pcd_merge = o3d.t.geometry.PointCloud(device)
    pcd_merge.point['positions'] = o3d.core.Tensor(points[..., :3], dtype, device)
    pcd_merge.point['intensity'] = o3d.core.Tensor(points[..., 3][:, None], dtype, device)
    o3d.t.io.write_point_cloud(pcd_file, pcd_merge)

def get_polygon(pc):
    points_2d = pc[:, :2]
    tri = Delaunay(points_2d)
    return tri



if __name__ == "__main__":
    root_dir = '/home/yaohan.lu/models/stereo-dsgn2/data/ww/training/velodyne'
    # tar_dir = f'/home/westwell/data/qomolo_labelled_20230620_{i}/training/ground_split'
    # if not os.path.exists(tar_dir):
    #     os.mkdir(tar_dir)
    files = os.listdir(root_dir)
    for file in files:
        pcd_path = os.path.join(root_dir, file)
        # bin_path = os.path.join(tar_dir, file)
        # if not os.path.exists(bin_path):
        pc = placeable_area_filter(pcd_path, min_height=0.1)
        polygon = get_polygon(pc)
